import random
import numpy as np
from contextlib import contextmanager
from omegaconf import OmegaConf
import dataclasses
from typing import List, Tuple
import torch
from verl import DataProto
from collections import defaultdict


@contextmanager
def all_seed(seed):
    """
    设置随机种子的上下文管理器，用于确保可重现性

    在进入上下文时保存当前随机状态，设置新的随机种子，
    在退出上下文时恢复原始随机状态

    参数:
        seed: 要设置的随机种子
    """
    random_state = random.getstate()
    np_random_state = np.random.get_state()

    try:
        random.seed(seed)
        np.random.seed(seed)
        yield
    finally:
        random.setstate(random_state)
        np.random.set_state(np_random_state)


def register_resolvers():
    """
    注册OmegaConf配置解析器

    添加自定义解析器函数:
    - mul: 乘法运算
    - int_div: 整数除法
    - not: 逻辑非

    这些解析器可在配置文件中使用，如${mul:x,y}
    """
    try:
        OmegaConf.register_new_resolver("mul", lambda x, y: x * y)
        OmegaConf.register_new_resolver(
            "int_div", lambda x, y: int(float(x) / float(y))
        )
        OmegaConf.register_new_resolver("not", lambda x: not x)
    except:
        pass  # 已经注册过，忽略异常


@dataclasses.dataclass
class GenerationsLogger:
    """
    生成结果记录器类

    用于将生成的样本记录到不同的日志系统中，如wandb和swanlab
    """

    def log(self, loggers, samples, step, _type="val"):
        """
        根据指定的日志器类型记录样本

        参数:
            loggers: 日志器类型列表，如['wandb', 'swanlab']
            samples: 要记录的样本列表
            step: 当前步数
            _type: 记录类型，默认为'val'(验证)
        """
        if "wandb" in loggers:
            self.log_generations_to_wandb(samples, step, _type)
        if "swanlab" in loggers:
            self.log_generations_to_swanlab(samples, step, _type)

    def log_generations_to_wandb(self, samples, step, _type="val"):
        """
        将样本记录到wandb表格中

        参数:
            samples: 要记录的样本列表，每个样本包含 (input, output, score, env_info)
            step: 当前步数
            _type: 记录类型，默认为'val'(验证)
        """
        import wandb

        # 处理样本数据，提取scenario_id用于列名
        processed_samples = []
        for i, sample in enumerate(samples):
            # 直接使用新格式：(input, output, score, env_info)
            input_text, output_text, score, env_info = sample[0], sample[1], sample[2], sample[3]
            scenario_id = env_info.get('scenario_id', f'{i+1}')
            env_config = env_info.get('env_config', '')
            # 在input前添加配置信息
            if env_config:
                input_text = f"配置信息: {env_config}\n\n{input_text}"
            
            processed_samples.append((scenario_id, input_text, output_text, score))

        # 为所有样本创建列名，使用scenario_id
        columns = ["step"] + sum(
            [
                [f"input_{sample[0]}", f"output_{sample[0]}", f"score_{sample[0]}"]
                for sample in processed_samples
            ],
            [],
        )

        if not hasattr(self, "table"):
            # 首次调用时初始化表格
            self.table = wandb.Table(columns=columns)

        # 创建具有相同列和现有数据的新表格
        # 这是为了解决 https://github.com/wandb/wandb/issues/2981#issuecomment-1997445737 中的问题
        new_table = wandb.Table(columns=columns, data=self.table.data)

        # 添加包含所有数据的新行
        row_data = []
        row_data.append(step)
        for scenario_id, input_text, output_text, score in processed_samples:
            row_data.extend([input_text, output_text, score])

        new_table.add_data(*row_data)

        # 更新引用并记录
        wandb.log({f"{_type}/generations": new_table}, step=step)
        self.table = new_table

    def log_generations_to_swanlab(self, samples, step, _type="val"):
        """
        将样本记录到swanlab文本中

        参数:
            samples: 要记录的样本列表，每个样本包含 (input, output, score, env_info)
            step: 当前步数
            _type: 记录类型，默认为'val'(验证)
        """
        import swanlab
        import os

        # 按tag分组样本
        tag_groups = defaultdict(list)
        
        for i, sample in enumerate(samples):
            # 直接使用新格式：(input, output, score, env_info)
            input_text, output_text, score, env_info = sample[0], sample[1], sample[2], sample[3]
            
            # 获取scenario_id和配置信息
            scenario_id = env_info.get('scenario_id', f'{i+1}')
            env_config = env_info.get('env_config', '')
            debug_info = env_info.get('debug_info', '')
            data_root = env_info.get('data_root', 'unknown data_root')
            
            # 从data_root路径中提取最后一层文件夹名称作为tag
            tag = os.path.basename(data_root.rstrip('/'))
            if not tag:  # 如果是根目录或空字符串
                tag = 'unknown'
            
            # 构建带配置信息的文本
            config_header = f"配置信息: {env_config}\n\n" if env_config else ""
            
            # 直接使用已格式化的调试文本
            debug_text = ""
            if debug_info and isinstance(debug_info, str):
                debug_text = debug_info
            
            row_text = f"""
            {config_header}input: {input_text}
            
            ---
            
            output: {output_text}
            
            ---
            
            score: {score}
            
            ---
            
            debug: 
            {debug_text}
            """
            
            # 使用scenario_id作为caption
            caption = f"sample {scenario_id} (score: {score})"
            swanlab_text = swanlab.Text(row_text, caption=caption)
            
            # 按tag分组
            tag_groups[tag].append(swanlab_text)

        # 为每个tag分别记录到swanlab
        for tag, swanlab_text_list in tag_groups.items():
            # 按照scenario_id的字符串顺序排序
            swanlab_text_list.sort(key=lambda x: x.caption.split('sample ')[1].split(' (score:')[0])
            
            # 将该tag下的样本按100条切分为小组，并使用group_idx标记
            max_group_size = 100
            for start_index in range(0, len(swanlab_text_list), max_group_size):
                end_index = start_index + max_group_size
                group_idx = start_index // max_group_size
                sub_group = swanlab_text_list[start_index:end_index]
                swanlab.log({f"{_type}/{tag}/group_{group_idx}/generations": sub_group}, step=step)
